package org.rbio.paper.methylglm.annotate;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

import org.broadinstitute.sting.gatk.datasources.reference.ReferenceDataSource;
import org.broadinstitute.sting.utils.codecs.refseq.RefSeqFeature;

// class to make TSS around the UCSC gene
public class TSSTable {
	// process
	public static void buildTSS(String ucsc_gene_file, String fa, int tss_up,
			int tss_down, String ofname) throws Exception {
		// initial gene table [ref [gene name [tx]]]
		Map<String, Map<String, List<RefSeqFeature>>> gene_ref_tab = UCSCGeneTable
				.read(ucsc_gene_file, fa);
		ReferenceDataSource fa_src = new ReferenceDataSource(new File(fa));

		// find TSS for each gene
		Map<String, Map<Integer, String>> tss_tab = new TreeMap<String, Map<Integer, String>>();
		for (String ref : gene_ref_tab.keySet()) {
			for (String gene_name : gene_ref_tab.get(ref).keySet()) {
				// strand of the gene and start
				char strand = '@';
				int start = -99999;
				for (RefSeqFeature tx : gene_ref_tab.get(ref).get(gene_name)) {
					// set the strand and start position the first time
					if (strand == '@') {
						if (tx.getStrand() == 1) {
							strand = '+';
							start = tx.getStart();
						} else {
							strand = '-';
							start = tx.getEnd();
						}
					}

					// compare with second or more tx
					if (strand == '+' && tx.getStart() < start) {
						start = tx.getStart();
					} else if (strand == '-' && tx.getEnd() > start) {
						start = tx.getEnd();
					}
				}
				int tss_start = start - tss_up;
				int tss_end = start + tss_down;
				if (strand == '-') {
					tss_start = start - tss_down;
					tss_end = start + tss_up;
				}

				// check for limits
				if (tss_start < 0) {
					tss_start = 1;
				}

				// check for chromosome end
				int ref_len = fa_src.getReference().getSequenceDictionary()
						.getSequence(ref).getSequenceLength();
				if (tss_end > ref_len) {
					tss_end = ref_len;
				}

				// TSS string
				String res = ref + ":" + tss_start + "-" + tss_end + "\t" + ref
						+ "\t" + gene_name + "\t" + strand + "\t" + tss_start
						+ "\t" + tss_end;

				// fill the tss table
				Map<Integer, String> ref_tab = tss_tab.get(ref);
				if (ref_tab == null) {
					ref_tab = new TreeMap<Integer, String>();
					tss_tab.put(ref, ref_tab);
				}
				ref_tab.put(tss_start, res);
			}
		}

		// write the tss_tab
		BufferedWriter obuf = new BufferedWriter(new FileWriter(ofname));
		obuf.write("HEADER\tref\tgene\tstrand\tstart." + tss_up + "\tend."
				+ tss_down + "\n");
		for (String ref : tss_tab.keySet()) {
			for (String rec : tss_tab.get(ref).values()) {
				obuf.write(rec + "\n");
			}
		}
		obuf.close();
	}
}
